import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from python_ai.common.xcommon import sep


def x_inverse_dummies(dummies_df, col_name):
    """
    Reversed operation of pandas.get_dummies

    https://stackoverflow.com/questions/50607740/reverse-a-get-dummies-encoding-in-pandas

    ----------------df == 1----------------
    blue  green
    0  False   True
    1  False   True
    2   True  False
    3   True  False
    4  False   True
    ----------------df[df == 1]----------------
       blue  green
    0   NaN    1.0
    1   NaN    1.0
    2   1.0    NaN
    3   1.0    NaN
    4   NaN    1.0
    ----------------df[df == 1].stack()----------------
    5
    MultiIndex([(0, 'green'),
                (1, 'green'),
                (2,  'blue'),
                (3,  'blue'),
                (4, 'green')],
               )
    0  green    1.0
    1  green    1.0
    2  blue     1.0
    3  blue     1.0
    4  green    1.0
    dtype: float64
    ----------------reset index----------------
    <class 'pandas.core.frame.DataFrame'>
       level_0 level_1    0
    0        0   green  1.0
    1        1   green  1.0
    2        2    blue  1.0
    3        3    blue  1.0
    4        4   green  1.0
    ----------------drop----------------
    <class 'pandas.core.frame.DataFrame'>
      level_1
    0   green
    1   green
    2    blue
    3    blue
    4   green
    ----------------rename----------------
    <class 'pandas.core.frame.DataFrame'>
       color
    0  green
    1  green
    2   blue
    3   blue
    4  green


    :param dummies_df:
    :param col_name:
    :return:
    """
    r = dummies_df[dummies_df == 1].stack().reset_index()
    cols = r.columns
    r.drop([cols[0], cols[-1]], axis=1, inplace=True)
    cols = r.columns
    r.rename(columns={cols[0]: col_name}, inplace=True)
    return r


def x_inverse_dummies_multiple(multi_dummies_df, *args):
    """
    Reversed operation of pandas.get_dummies

    https://stackoverflow.com/questions/50607740/reverse-a-get-dummies-encoding-in-pandas

    :param multi_dummies_df:
    :param args:
    :return:
    """
    n_cols = len(multi_dummies_df.columns)
    n_groups = len(args)
    result_list = []
    for k, v in enumerate(args):
        start_idx, col_name = v
        if k == n_groups - 1:
            end_idx = n_cols
        else:
            end_idx = args[k + 1][0]
        dummies_df = multi_dummies_df.iloc[:, start_idx:end_idx]
        rv = x_inverse_dummies(dummies_df, col_name)
        len_col = len(col_name)
        rv[rv.columns[0]] = rv[rv.columns[0]].map(lambda x: x[len_col+1:])
        result_list.append(rv)
    r = pd.concat(result_list, axis=1)
    return r



if '__main__' == __name__:
    df = pd.DataFrame([['green', 'a'],
                       ['green', 'b'],
                       ['blue', 'a'],
                       ['blue', 'b'],
                       ['green', 'a'],
                       ['red', 'd']],
                      columns=['color', 'letter'])
    print(df)

    sep('onehot')
    # get_dummies is just onehot encoder
    df1 = pd.get_dummies(df)
    print(type(df1))
    print(df1)
    df1rv = x_inverse_dummies_multiple(df1, [0, 'color'], [3, 'letter'])
    print(df1rv)
    print('EQUALS', df.equals(df1rv))

    sep('onehot on col0')
    df2 = pd.get_dummies(df[df.columns[0]])
    print(df2)
    df2rv = x_inverse_dummies(df2, 'color')
    print(type(df2rv))
    print(df2rv)
